import os.path
import random
from torch.utils.data import Dataset
import torch
import numpy as np
from plyfile import PlyData, PlyElement
from sklearn.neighbors import KDTree

def get_data(path, item):

    node_pos = np.load(f"{path}/Testset_track_B/Inference/centroid_{item}.npy").astype(np.float32)
    areas = np.load(f"{path}/Testset_track_B/Auxiliary/area_{item}.npy").astype(np.float32)

    return node_pos, areas.reshape(-1,1)

class Car_Dataset(Dataset):
    def __init__(self,
                 data_path,
                 mode="test",
                 adj_num = 3
                 ):

        super(Car_Dataset, self).__init__()
        self.dataloc = []

        self.mode = mode
        self.adj_num = adj_num

        self.fn = data_path

        file_list = os.listdir(f"{self.fn}/Testset_track_B/Inference")
        for file_name in file_list:
            if file_name.startswith('centroid'):
                split_string = file_name.split('.')
                split_string1 = split_string[0].split('_')
                self.dataloc.append(split_string1[1])

    def __len__(self):
        return len(self.dataloc)

    def get_singel(self, item):

        node_pos, areas = get_data(
                        self.fn,
                        self.dataloc[item],
                        )

        if self.adj_num>0:
            edges = self.knn_edges(node_pos, self.adj_num)

        node_pos = torch.from_numpy(node_pos).float()
        node_pos = self.scale_pos(node_pos)

        areas = torch.from_numpy(areas).float()
        areas = self.scale_area(areas)

        input = {'node_pos': node_pos,
                'edges': edges,
                'areas': areas
                 }

        return input, self.dataloc[item]
    
    def knn_edges(self, node_pos, k):
        node_pos_np = node_pos.detach().numpy()
        # 构建KD树
        tree = KDTree(node_pos_np)
        # 查询每个点的最近k个邻居，返回距离和索引
        dists, indices = tree.query(node_pos_np, k=k+1)  # k+1是因为最近的一个邻居是点本身
        # 删除每个点自身的索引
        indices = indices[:, 1:]  # 第一列是每个点自身，所以去掉
        # 构建边的索引
        num_points = node_pos.shape[0]
        i_indices = torch.arange(num_points).view(-1, 1).repeat(1, k)
        edges = torch.stack([i_indices.flatten(), torch.from_numpy(indices.flatten())], dim=1)

        return edges.long()

    def scale_area(self, area):

        area_min = 3.0709e-09
        area_max = 0.00051785

        area_01 =  (area - area_min) / (area_max - area_min)

        return area_01

    def scale_pos(self, pos):

        pos_min = torch.tensor([-1.1514,-1.0214, 0.0019]).to(pos.device)
        pos_max = torch.tensor([4.0932, 1.0214, 1.7619]).to(pos.device)

        node_pos = (pos - pos_min.reshape(-1,3)) / (pos_max.reshape(-1,3) - pos_min.reshape(-1,3))

        return node_pos

    def __getitem__(self, item):

        input = self.get_singel(item)
        return input